import os
import time

from selenium import webdriver
from selenium.webdriver.chrome.options import Options

from config.app_config import AppConfig
from src.custom_logger import CustomLogger


class SeleniumManager:
    def __init__(self):
        self.options = Options()
        self.options.headless = AppConfig.HEADLESS
        self.options.add_argument("--log-level=3")
        self.options.add_argument(f"user-data-dir={AppConfig.PROFILE_PATH}") # 使用已有的Chrome配置
        self.driver = webdriver.Chrome(executable_path=AppConfig.CHROMEDRIVER_PATH, options=self.options)
        self.logger = CustomLogger()

    def manual_login(self, login_url="https://leetcode.cn/accounts/login/?next=%2F"):
        """Manual Login"""
        self.logger.info("Opening the login page, please manually enter your username and password to log in.")
        self.driver.get(login_url)

        # Wait for the user to manually log in and then input 'y' or 'yes' to continue
        # confirm = input("If you have successfully logged in, please enter 'Y' or 'yes' to continue, or press Enter to skip login: ").strip().lower()
        confirm = 'y'
        time.sleep(5)
        if confirm in ['y', 'yes']:
            self.logger.info("Continuing with the subsequent operations.")
        else:
            self.logger.info("Skipping login and continuing with the subsequent operations.")

    def get_cookies(self):
        return {cookie['name']: cookie['value'] for cookie in self.driver.get_cookies()}
    
    def get_csrf_token(self):
        """Retrieve the CSRF token from cookies."""
        cookies = self.get_cookies()
        return cookies.get('csrftoken', None)
    

    def get_headers(self):
        """Generate a probable set of headers for HTTP requests.

        Note: Selenium WebDriver does not support directly retrieving the headers
        of the current browser session's HTTP requests. This method provides a way
        to construct a probable set of request headers based on the current session
        cookies and CSRF token (if available).
        """
        cookies = self.get_cookies()
        csrf_token = self.get_csrf_token()

        # Constructing a probable set of headers
        headers = {
            "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/119.0.0.0 Safari/537.36",
            "Accept-Encoding": "gzip, deflate",
            "Accept-Language": "zh-CN,zh;q=0.9",
            # "Accept-Language": "en-US,en;q=0.9",
            "Accept": "*/*",
            "Cookie": "; ".join([f"{name}={value}" for name, value in cookies.items()]),
            "Origin": "https://leetcode.cn",
        }

        if csrf_token:
            headers["X-Csrftoken"] = csrf_token

        return headers
    
    def close(self):
        """Close WebDriver"""
        self.driver.quit()

if __name__ == "__main__":
    # 使用示例
    selenium_chrome = SeleniumChrome()
    cookies = selenium_chrome.login_and_get_cookies()
    selenium_chrome.close()
